import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs
import platform

# Set font based on operating system
if platform.system() == 'Windows':
    plt.rcParams['font.sans-serif'] = ['SimHei']  # Windows common Chinese font
elif platform.system() == 'Darwin':  # macOS
    plt.rcParams['font.sans-serif'] = ['STHeiti']
else:  # Linux or other
    plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']  # Example font, may vary based on system

# Set minus sign to display correctly
plt.rcParams['axes.unicode_minus'] = False


# Generate synthetic data with higher cluster standard deviation for more variability
def generate_data(n_samples=500, n_clusters=3, cluster_std=1.5):
    X, _ = make_blobs(n_samples=n_samples, centers=n_clusters, cluster_std=cluster_std)
    return X

# Initialize centroids randomly from data points
def initialize_centroids(X, k):
    np.random.seed(42)
    random_indices = np.random.choice(X.shape[0], k, replace=False)
    return X[random_indices]

# Assign each point to the nearest centroid
def assign_clusters(X, centroids):
    distances = np.sqrt(((X - centroids[:, np.newaxis])**2).sum(axis=2))
    return np.argmin(distances, axis=0)

# Update centroids by moving them gradually toward the mean position of assigned points
def gradual_update_centroids(X, labels, centroids, k, step_size=0.2):
    new_centroids = np.array([X[labels == i].mean(axis=0) if np.any(labels == i) else centroids[i] for i in range(k)])
    centroids = centroids + step_size * (new_centroids - centroids)  # Move centroids toward their target position
    return centroids

# Animate the clustering process with gradual centroid updates
def animate_clustering(X, k, max_iterations=15, step_size=0.2):
    centroids = initialize_centroids(X, k)
    labels = np.zeros(X.shape[0])

    plt.ion()  # Turn on interactive mode for animation
    fig, ax = plt.subplots(figsize=(8, 6))
    
    for iteration in range(max_iterations):
        # Step 1: Assign clusters
        labels = assign_clusters(X, centroids)
        
        # Clear the plot and show clusters
        ax.cla()
        
        # Plot points with cluster colors
        for i in range(k):
            cluster_points = X[labels == i]
            ax.scatter(cluster_points[:, 0], cluster_points[:, 1], label=f'簇 {i + 1}')
        
        # Plot centroids
        ax.scatter(centroids[:, 0], centroids[:, 1], s=200, c='black', marker='X', label='中心点')
        ax.set_title(f"K-Means 聚类 - 第 {iteration + 1} 次迭代")
        ax.set_xlabel("特征 1")
        ax.set_ylabel("特征 2")
        ax.legend()

        plt.draw()
        plt.pause(1)  # Pause to create animation effect

        # Step 2: Gradual update of centroids
        new_centroids = gradual_update_centroids(X, labels, centroids, k, step_size)
        
        # Check for convergence (if centroids have barely moved)
        if np.allclose(centroids, new_centroids, atol=1e-2):
            print("已收敛")
            break
        centroids = new_centroids

    plt.ioff()  # Turn off interactive mode
    plt.show()

if __name__ == "__main__":
    # Generate data and animate K-Means clustering with gradual updates
    X = generate_data(n_samples=500, n_clusters=3, cluster_std=1.5)
    k = 3  # Number of clusters
    animate_clustering(X, k)